import time
import rclpy
import threading
from utils.moveit2 import *
from utils.robot_urdf import *
from utils.specific_pub_node_move_arm import *
from utils.particle_swarm_optimization import *
from rclpy.executors import MultiThreadedExecutor

# 将机械臂参数写入urdf文件
def write_robot_urdf(name, link_number, link_width_list, link_length_list, joint_type_list, axis_list, output_path):
    robot = Robot(name)
    link_name_list = [f'link{i}' for i in range(link_number + 1)]
    joint_name_list = [f'joint{i}' for i in range(link_number)]
    
    for i in range(link_number):
        if i < link_number - 1:
            robot.add_link(link_name_list[i], link_name_list[i+1], link_width_list[i], link_length_list[i],
                           joint_name_list[i], joint_type_list[i], axis_list[i])
        else:
            robot.add_link(link_name_list[i], link_name_list[i+1], link_width_list[i], link_length_list[i],
                           joint_name_list[i], joint_type_list[i], axis_list[i], 1)

    robot.output(output_path)

# 粒子群参数定义
dimension = 7
swarm_size = 5
iteration_size = 10
low_limit = 0.1
high_limit = 1.5
min_lims = np.array([low_limit] * dimension)
max_lims = np.array([high_limit] * dimension)

particle_swarm = Swarm(dimension, swarm_size, iteration_size, min_lims, max_lims)

# 机械臂其余参数定义
link_number = 7
joint_type = ['continuous' for _ in range(link_number)]
joint_axis = ['z', 'x', 'y', 'z', 'x', 'z', 'x']
doc_path = "./src/mybot_description/urdf/try.urdf"
link_width = [0.3, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2]

#主要程序循环
for i in range(0, iteration_size):
    for j in range(0, swarm_size):
        # 获取粒子生成的机械臂参数并写入urdf
        particle_number, link_length = particle_swarm.get_particle_position()
        write_robot_urdf('robot', link_number, link_width, link_length, joint_type, joint_axis, doc_path)

        # 启动moveit进程
        moveit_process = run_moveit2()
        rclpy.init()

        # 发布与订阅节点的设置
        executor_plan_command = MultiThreadedExecutor()
        executor_move_arm_command = MultiThreadedExecutor()

        pub_node_plan_command = Publisher("pub_node_plan_command", Empty, "/rviz/moveit/plan", Empty(), 2)
        pub_node_move_arm_command = MoveArmPublisher("pub_node_move_arm_command", PoseStamped, "/rviz/moveit/move_marker/goal_link6")
        sub_node_get_path_info = Subscriber("sub_node_get_path_info", DisplayTrajectory, "/display_planned_path")
        print("nodes initialization succeeded.")

        stop_signal = True

        # 子线程执行内容
        def plan_command():
            while(stop_signal):
                time.sleep(2)
                rclpy.spin_once(pub_node_plan_command, executor=executor_plan_command)

        def move_arm_command():
            rclpy.spin_once(pub_node_move_arm_command, executor=executor_move_arm_command)
        

        thread = threading.Thread(target=plan_command)
        thread.start()

        # 确认moveit已启动
        rclpy.spin_once(sub_node_get_path_info)
        sub_node_get_path_info.destroy_node()

        stop_signal = False
        thread.join()
        time.sleep(5)
        print("moveit and rviz start succeeded")
        

        #设定目标点
        pub_node_move_arm_command.set_goal(x=0.95, y=-0.37, z=1.7)
        thread = threading.Thread(target=move_arm_command)
        thread.start()

        time.sleep(5)
        sub_node_get_path_info = Subscriber("sub_node_get_path_info", DisplayTrajectory, "/display_planned_path")
        rclpy.spin_once(pub_node_plan_command, executor= executor_plan_command)
        rclpy.spin_once(sub_node_get_path_info)

        time.sleep(5)
        # 获取运动轨迹时间
        time_from_start_sec = sub_node_get_path_info.message.trajectory[0].joint_trajectory.points[-1].time_from_start.sec
        time_from_start_nsec = sub_node_get_path_info.message.trajectory[0].joint_trajectory.points[-1].time_from_start.nanosec

        pub_node_plan_command.destroy_node()
        pub_node_move_arm_command.destroy_node()
        sub_node_get_path_info.destroy_node()

        # 更新单个粒子的代价
        particle_swarm.return_particle_best(particle_number, time_from_start_sec + time_from_start_nsec * 1e-9)
        print("the particle %d of iteration %d's cost is: %.4f seconds"%(j, i, time_from_start_sec + time_from_start_nsec * 1e-9))

        rclpy.shutdown()
        kill_moveit2(moveit_process)
        print('the particle %d of iteration %d finished.'%(j, i))

    # 一次迭代后更新粒子群
    particle_swarm.update_swarm()
    position, cost = particle_swarm.get_swarm_best()
    position = [round(x, 3) for x in position]
    print("best position:", position, "best cost: {: .7f}".format(cost))

# 全部完成后的最优
position, cost = particle_swarm.get_swarm_best()
position = [round(x, 3) for x in position]
print("global best position:", position, "global best cost: {: .7f}".format(cost))

